import re
import matplotlib.pyplot as plt
import os

def parse_log(log_file):
    epochs = []
    losses = []
    accuracies = []
    mAPs = []
    rank1s = []

    with open(log_file, 'r') as f:
        for line in f:
            match = re.search(r'Epoch\[(\d+)\] Iteration\[\d+/(\d+)\] Loss: ([\d.]+), Acc: ([\d.]+), Base Lr: ([\d.e-]+)', line)
            if match:
                epoch = int(match.group(1))
                total_iters = int(match.group(2))
                loss = float(match.group(3))
                acc = float(match.group(4))
                if epoch not in epochs:
                    epochs.append(epoch)
                    losses.append(loss)
                    accuracies.append(acc)
            match_mAP = re.search(r'Validation Results - Epoch: (\d+).*mAP: ([\d.]+)%', line)
            if match_mAP:
                epoch = int(match_mAP.group(1))
                mAP = float(match_mAP.group(2)) / 100
                mAPs.append((epoch, mAP))
            match_rank1 = re.search(r'CMC curve, Rank-1\s+:([\d.]+)%', line)
            if match_rank1:
                rank1 = float(match_rank1.group(1)) / 100
                rank1s.append((epoch, rank1))

    return epochs, losses, accuracies, mAPs, rank1s

def plot_training(log_file):
    epochs, losses, accuracies, mAPs, rank1s = parse_log(log_file)

    plt.figure(figsize=(12, 8))
    plt.subplot(2, 2, 1)
    plt.plot(epochs, losses, label='Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss')
    plt.legend()

    plt.subplot(2, 2, 2)
    plt.plot(epochs, accuracies, label='Training Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title('Training Accuracy')
    plt.legend()

    if mAPs:
        mAP_epochs, mAP_values = zip(*mAPs)
        plt.subplot(2, 2, 3)
        plt.plot(mAP_epochs, mAP_values, label='mAP')
        plt.xlabel('Epoch')
        plt.ylabel('mAP')
        plt.title('Validation mAP')
        plt.legend()

    if rank1s:
        rank1_epochs, rank1_values = zip(*rank1s)
        plt.subplot(2, 2, 4)
        plt.plot(rank1_epochs, rank1_values, label='Rank-1')
        plt.xlabel('Epoch')
        plt.ylabel('Rank-1')
        plt.title('Validation Rank-1')
        plt.legend()

    plt.tight_layout()
    plt.savefig('training_plots.png')
    plt.show()

if __name__ == '__main__':
    log_file = r"C:\Users\xuboyang\Desktop\TransReID-main\logs\0321_market_vit_base\train_log.txt"
    plot_training(log_file)